import hashlib
import json
import os
import time

import torch

from util.drawing_util import get_forecast_comparison_chart_numpy
from util.file_util import find_last_file


def save_model(model,optimizer,epoch,loss,lstm_config):
    try:
        checkpoint = {
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'epoch': epoch,
            'lstm_config': lstm_config,
            'loss': loss.item()
        }
        os.makedirs('checkpoint', exist_ok=True)
        lstm_config_md5 = lstm_config_to_md5(lstm_config)
        # 一小时保存一次 避免保存文件过多 保存到项目的checkpoint目录中
        torch.save(checkpoint,
                   os.path.join('checkpoint', f'checkpoint_{lstm_config_md5}_{time.strftime("%Y-%m-%d_%H", time.localtime())}.pth'))
    except Exception as e:
        print(e)


def load_checkpoint_last_file(checkpoint_directory_path,lstm_config):
    try:
        lstm_config_md5 = lstm_config_to_md5(lstm_config)
        checkpoint_last_file = find_last_file(checkpoint_directory_path,lstm_config_md5)
        if checkpoint_last_file is not None and os.path.exists(checkpoint_last_file):
            checkpoint = torch.load(checkpoint_last_file)
            return checkpoint
    except Exception as e:
        print(e)


def summary_add_image(writer,x_batch,outputs,y_batch_all,global_step):
    try:
        forecast_comparison_chart_img = get_forecast_comparison_chart_numpy(x_batch[:, -20:].detach().cpu().numpy(),
                                                                            outputs.detach().cpu().numpy(),
                                                                            y_batch_all.detach().cpu().numpy())
        writer.add_image("forecast_comparison_chart_img", forecast_comparison_chart_img, global_step, dataformats="HWC")
    except Exception as e:
        print(e)

def lstm_config_to_md5(lstm_config):
    lstm_config_md5 = "md5"
    if lstm_config is not None:
        lstm_config_src = json.dumps(lstm_config)
        lstm_config_md5 = hashlib.md5(lstm_config_src.encode()).hexdigest()
    return lstm_config_md5